#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2025/5/30 16:54
@Author  : HZP
@File    : 1.摘要缓冲混合记忆示例.py
"""
from dotenv import load_dotenv
from langchain.memory import ConversationSummaryBufferMemory
from langchain_community.chat_message_histories import FileChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import MessagesPlaceholder,ChatPromptTemplate
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_openai import OpenAI, ChatOpenAI

load_dotenv()

store = {}

def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = FileChatMessageHistory(f"./chat_history_{session_id}.txt")
    return store[session_id]

#2、构建摘要器
summary = ConversationSummaryBufferMemory(
    llm=ChatOpenAI(temperature=0.6, model="qwq"),
    input_key="query",
    return_messages=True,
    max_token_limit=1000
)

# 1、prompt提示词
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是中电三公司聊天机器人,根据上下文信息回答用户问题"),
    MessagesPlaceholder("history"),
    ("human", "{query}")
])
# 2、创建模型
llm = ChatOpenAI(model="deepseek-r1:70b", temperature=0.6)

#3、构建链
chain=prompt|llm|StrOutputParser()

# 2、摘要缓冲混合记忆
merry =RunnableWithMessageHistory(
    chain,
    get_session_history,
    input_messages_key="query",
    history_messages_key="history"
)

while True:
    query = input("Human:")
    if query == "q":
        break
    chain_input = {"query": query, "language": "中文"}
    chain_output = merry.stream(
        chain_input,
        config={"configurable":{"session_id": "123"}}
    )
    print("AI:", flush=True, end="")
    for chunk in chain_output:
        if not chunk:
            break
        print(chunk, flush=True, end="")
    print("", flush=False, end="")